/**********************************************************************
* file:   PacketFilter.c
* date:   Oct 13 2010
* Author: Liting
*
* Description:parse src and dst ip packet for each VM
*
* Compile with:
* gcc -Wall -pedantic PacketFilter.c -lpcap
*
* Usage:
* a.out (interface) (ip_address)
*
**********************************************************************/

#include <pcap.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netinet/if_ether.h>
#include <net/ethernet.h>
#include <netinet/ether.h>
#include <netinet/ip.h>
#include <signal.h>

#ifndef ETHER_HDRLEN
#define ETHER_HDRLEN 14
#endif

#ifndef TOP_K
#define TOP_K 57
#endif

#ifndef TIMEOUT
#define TIMEOUT 1500
#endif

struct in_addr src_addr;		/* vm ip					 */
uint32_t dst_addr[TOP_K] = {0};	/* dst vm					 */
uint32_t dst_count[TOP_K] = {0};
static uint32_t packet_out = 0;
static int index = 0;
pcap_t* descr;

u_int16_t handle_ethernet
        (u_char *args,const struct pcap_pkthdr* pkthdr,const u_char*
        packet);
u_char* handle_IP
        (u_char *args,const struct pcap_pkthdr* pkthdr,const u_char*
        packet);

/*
 * Structure of an internet header, naked of options.
 *
 * We declare ip_len and ip_off to be short, rather than u_short
 * pragmatically since otherwise unsigned comparisons can result
 * against negative integers quite easily, and fail in subtle ways.
 */
struct my_ip {
	u_int8_t	ip_vhl;		/* header length, version */
#define IP_V(ip)	(((ip)->ip_vhl & 0xf0) >> 4)
#define IP_HL(ip)	((ip)->ip_vhl & 0x0f)
	u_int8_t	ip_tos;		/* type of service */
	u_int16_t	ip_len;		/* total length */
	u_int16_t	ip_id;		/* identification */
	u_int16_t	ip_off;		/* fragment offset field */
#define	IP_DF 0x4000			/* dont fragment flag */
#define	IP_MF 0x2000			/* more fragments flag */
#define	IP_OFFMASK 0x1fff		/* mask for fragmenting bits */
	u_int8_t	ip_ttl;		/* time to live */
	u_int8_t	ip_p;		/* protocol */
	u_int16_t	ip_sum;		/* checksum */
	struct	in_addr ip_src,ip_dst;	/* source and dest address */
};

/* looking at ethernet headers */
void my_callback(u_char *args,const struct pcap_pkthdr* pkthdr,const u_char*
        packet)
{
    u_int16_t type = handle_ethernet(args,pkthdr,packet);
    if(type == ETHERTYPE_IP)
    {/* handle IP packet */
        handle_IP(args,pkthdr,packet);
    }else if(type == ETHERTYPE_ARP)
    {/* handle arp packet */
    }
    else if(type == ETHERTYPE_REVARP)
    {/* handle reverse arp packet */
    }
}

int linearSearch( const uint32_t array[], uint32_t key, int size )
{
   int i;
   for (i = 0; i < size; ++i) {
      if (array[i] == key) {
         return i;
      }
   }
   return -1;
}

u_char* handle_IP
        (u_char *args,const struct pcap_pkthdr* pkthdr,const u_char*
        packet)
{
    const struct my_ip* ip;
    u_int length = pkthdr->len;
    u_int hlen,off,version;

    int len;

    /* jump pass the ethernet header */
    ip = (struct my_ip*)(packet + sizeof(struct ether_header));
    length -= sizeof(struct ether_header);

    /* check to see we have a packet of valid length */
    if (length < sizeof(struct my_ip))
    {
        printf("truncated ip %d",length);
        return NULL;
    }

    len     = ntohs(ip->ip_len);
    hlen    = IP_HL(ip); /* header length */
    version = IP_V(ip);/* ip version */

    /* check version */
    if(version != 4)
    {
      fprintf(stdout,"Unknown version %d\n",version);
      return NULL;
    }

    /* check header length */
    if(hlen < 5 )
    {
        fprintf(stdout,"bad-hlen %d \n",hlen);
    }

    /* see if we have as much packet as we should */
    if(length < len)
        printf("\ntruncated IP - %d bytes missing\n",len - length);

    /* Check to see if we have the first fragment */
    off = ntohs(ip->ip_off);
    if((off & 0x1fff) == 0 )/* aka no 1's in first 13 bits */
    {
/* print SOURCE DESTINATION */
//        fprintf(stdout,"IP: ");
//        fprintf(stdout,"src: %s \n",
//        		inet_ntoa(ip->ip_src));
        if(ip->ip_src.s_addr == src_addr.s_addr && index < TOP_K){
		packet_out++;
//        fprintf(stdout,"dst: %s  packet_out: %u\n",
//               inet_ntoa(ip->ip_dst), packet_out);

        int rt = linearSearch(dst_addr, ip->ip_dst.s_addr, TOP_K);
        if(rt == -1){
        	dst_addr[index] = ip->ip_dst.s_addr;
        	dst_count[index]++;
			index++;
        }
        if(rt != -1){
			dst_count[rt]++;
        }
        if(rt > TOP_K-1)
			fprintf(stderr, " linearSearch returns wrong index ");

//        fprintf(stdout, "dst_addr %u %u %u %u %u %u %u %u %u %u\n", dst_addr[0], dst_addr[1], dst_addr[2], dst_addr[3],
//																	dst_addr[4], dst_addr[5], dst_addr[6], dst_addr[7],
//																	dst_addr[8], dst_addr[9]);
//        fprintf(stdout, "dst_count %u %u %u %u %u %u %u %u %u %u\n", dst_count[0], dst_count[1], dst_count[2], dst_count[3],
//																	dst_count[4], dst_count[5], dst_count[6], dst_count[7],
//																	dst_count[8], dst_count[9]);

        }
//        dst_buffer[0] = ip->ip_dst.s_addr;
//        fprintf(stdout,"dst_buffer[0] = %u\n", dst_buffer[0]);
    }

    return NULL;
}

/* handle ethernet packets */

u_int16_t handle_ethernet
        (u_char *args,const struct pcap_pkthdr* pkthdr,const u_char*
        packet)
{
    u_int caplen = pkthdr->caplen;
    struct ether_header *eptr;  /* net/ethernet.h */
    u_short ether_type;

    if (caplen < ETHER_HDRLEN)
    {
        fprintf(stdout,"Packet length less than ethernet header length\n");
        return -1;
    }

    /* lets start with the ether header... */
    eptr = (struct ether_header *) packet;
    ether_type = ntohs(eptr->ether_type);
    return ether_type;
}

void terminate_process(int signum)
{
   pcap_breakloop(descr);
   pcap_close(descr);
}

int main(int argc,char **argv)
{
    char *dev;
    char errbuf[PCAP_ERRBUF_SIZE];

    bpf_u_int32 maskp;          /* subnet mask               */
    bpf_u_int32 netp;           /* ip                        */
    u_char* args = NULL;
    FILE* fp;

    if(argc < 3){
        fprintf(stdout,"Usage: %s \"interface\" + \"ip address\" + \"options\"\n", argv[0]);
        return 0;
    }

    /* grab a device to peak into... */
    dev = pcap_lookupdev(errbuf);
    if(dev == NULL)
    { printf("%s\n",errbuf); exit(1); }

    dev = argv[1];
    /* ask pcap for the network address and mask of the device */
    pcap_lookupnet(dev,&netp,&maskp,errbuf);

    /* open device for reading. NOTE: defaulting to
     * promiscuous mode*/
    descr = pcap_open_live(dev,BUFSIZ,1,-1,errbuf);
    if(descr == NULL)
    { printf("pcap_open_live(): %s\n",errbuf); exit(1); }

    inet_aton(argv[2], &src_addr);
    fprintf(stdout, "PacketFilter: \"interface\" %s + \"ip address\" %s \n", argv[1], argv[2]);

    signal(SIGALRM, terminate_process);
    alarm(100);
    /* ... and loop */
    pcap_loop(descr,-1,my_callback,args);

    int i;
    char output[100];
    strcpy(output,"./");
    strcat(output, argv[1]);
    //strcat(output, ".txt");
    fprintf(stdout, "\nwriting to ");
    fprintf(stdout, output);

    if((fp = fopen(output, "w")) != (FILE *)0) {
    	fprintf(fp,"src: %u \n", src_addr.s_addr);
    	fprintf(stdout, "\nsrc: %u \n", src_addr.s_addr);

    	fprintf(fp,"count: ");
    	fprintf(stdout,"count: ");

    	for(i = 0; i < TOP_K; i++) {
    	fprintf(fp, "%u ", dst_count[i]);
    	fprintf(stdout, "%u ", dst_count[i]);
    	}

    	fprintf(fp,"\ndst: ");
    	fprintf(stdout,"\ndst: ");

    	for(i = 0;i < TOP_K; i++) {
    	fprintf(fp, "%u	", dst_addr[i]);
    	fprintf(stdout, "%u	", dst_addr[i]);
    	}


    } else {
    	fprintf(stderr, "File open error");
    }
    fclose(fp);

    fprintf(stdout,"\nfinished\n");
    return 0;
}
